use shared_memory::*;
use zmq::{Context as ZmqContext, REP};
use serde::{Deserialize, Serialize};
use image::DynamicImage;
use std::io::Cursor;
use anyhow::Result;

#[derive(Serialize, Deserialize, Debug)]
struct ControlMsg {
    command: String,
    shmem_id: String,
    width: u32,
    height: u32,
    params: Option<serde_json::Value>,
}

#[derive(Serialize, Deserialize)]
struct ResponseMsg {
    status: String,
    result_shmem_id: Option<String>,
    error: Option<String>,
}

fn process_image(
    input: &[u8],
    _width: u32,
    _height: u32,
    command: &str
) -> Result<Vec<u8>> {
    let img = image::load_from_memory(input)?;
    let result = match command {
        "grayscale" => img.grayscale().to_rgba8(),
        "invert" => {
            let mut img = img.to_rgba8();
            image::imageops::invert(&mut img);
            img
        }
        "blur" => img.blur(5.0).to_rgba8(),
        "rotate90" => img.rotate90().to_rgba8(),
        "canny_edge" => {
            // 实现简单的Sobel边缘检测
            let gray_img = img.to_luma8();
            let mut edges = gray_img.clone();
            
            for y in 1..gray_img.height()-1 {
                for x in 1..gray_img.width()-1 {
                    let mut gx = 0i16;
                    let mut gy = 0i16;
                    
                    for dy in -1..=1 {
                        for dx in -1..=1 {
                            let pixel = gray_img.get_pixel((x as i32 + dx) as u32, (y as i32 + dy) as u32)[0] as i16;
                            
                            if dx == -1 {
                                gx -= pixel;
                            } else if dx == 1 {
                                gx += pixel;
                            }
                            
                            if dy == -1 {
                                gy -= pixel;
                            } else if dy == 1 {
                                gy += pixel;
                            }
                        }
                    }
                    
                    let magnitude = (gx.abs() + gy.abs()).min(255) as u8;
                    edges.put_pixel(x, y, image::Luma([magnitude]));
                }
            }
            
            DynamicImage::ImageLuma8(edges).to_rgba8()
        }
        _ => return Err(anyhow::anyhow!("Unknown command: {}", command)),
    };

    let mut buffer = Vec::new();
    result.write_to(&mut Cursor::new(&mut buffer), image::ImageOutputFormat::Png)?;
    Ok(buffer)
}

fn main() -> Result<()> {
    let ctx = ZmqContext::new();
    let socket = ctx.socket(REP)?;
    
    // 使用临时文件路径确保跨平台兼容性
    #[cfg(not(target_os = "windows"))]
    let ipc_path = "ipc:///tmp/image_processor.ipc";
    #[cfg(target_os = "windows")]
    let ipc_path = "tcp://127.0.0.1:5555";
    
    socket.bind(ipc_path)?;
    println!("Rust processor started. Listening on: {}", ipc_path);

    // 存储共享内存对象防止提前释放
    let mut shmem_storage: Vec<shared_memory::Shmem> = Vec::new();

    loop {
        let msg_str = match socket.recv_string(0) {
            Ok(Ok(msg)) => msg,
            Ok(Err(e)) => {
                eprintln!("Invalid UTF-8 message: {:?}", e);
                continue;
            }
            Err(e) => {
                eprintln!("Error receiving message: {}", e);
                continue;
            }
        };
        
        let control: ControlMsg = match serde_json::from_str(&msg_str) {
            Ok(c) => c,
            Err(e) => {
                let response = ResponseMsg {
                    status: "error".to_string(),
                    result_shmem_id: None,
                    error: Some(format!("Invalid control message: {}", e)),
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
                continue;
            }
        };

        // 访问共享内存
        let shmem = match ShmemConf::new().os_id(&control.shmem_id).open() {
            Ok(m) => m,
            Err(e) => {
                let response = ResponseMsg {
                    status: "error".to_string(),
                    result_shmem_id: None,
                    error: Some(format!("Shmem open failed: {}", e)),
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
                continue;
            }
        };

        // 处理图像
        let result = process_image(
            unsafe { std::slice::from_raw_parts(shmem.as_ptr(), shmem.len()) },
            control.width,
            control.height,
            &control.command
        );

        match result {
            Ok(output_buffer) => {
                // 创建结果共享内存
                let result_id = format!("{}-result", control.shmem_id);
                let result_shmem = match ShmemConf::new()
                    .size(output_buffer.len())
                    .os_id(&result_id)
                    .create()
                {
                    Ok(s) => s,
                    Err(e) => {
                        let response = ResponseMsg {
                            status: "error".to_string(),
                            result_shmem_id: None,
                            error: Some(format!("Result shmem creation failed: {}", e)),
                        };
                        socket.send(&serde_json::to_string(&response)?, 0)?;
                        continue;
                    }
                };

                unsafe {
                    std::ptr::copy_nonoverlapping(
                        output_buffer.as_ptr(),
                        result_shmem.as_ptr(),
                        output_buffer.len()
                    );
                }

                // 存储共享内存防止提前释放
                shmem_storage.push(result_shmem);
                
                // 清理旧共享内存（保留最近5个）
                if shmem_storage.len() > 5 {
                    shmem_storage.remove(0);
                }

                let response = ResponseMsg {
                    status: "success".to_string(),
                    result_shmem_id: Some(result_id),
                    error: None,
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
            }
            Err(e) => {
                let response = ResponseMsg {
                    status: "error".to_string(),
                    result_shmem_id: None,
                    error: Some(e.to_string()),
                };
                socket.send(&serde_json::to_string(&response)?, 0)?;
            }
        }
    }
}